package com.fastjars.mybatis.dynamic;

import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.lang.Nullable;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

/**
 * DataSource
 * 动态数据源实现读写分离
 * @author Brant Liu <br> 邮箱：<br>lbf1988@qq.com <br>日期：<br>2018/10/10
 * @version 1.0.0
 */
public class DataSource extends AbstractRoutingDataSource {
    /**
     * 主库，用于写数据
     */
    private Object masterDataSource;
    /**
     * 从库，用于读数据
     */
    private List<Object> slaveDataSources;

    /**
     * 读数据源数量
     */
    private int slaveDataSourceSize;

    /**
     * 获取读数据源方式，0：随机，1：轮询
     */
    private int readDataSourcePollPattern = 0;

    private AtomicLong counter = new AtomicLong(0);

    private static final Long MAX_POOL = Long.MAX_VALUE;

    private final Lock lock = new ReentrantLock();

    @Override
    public void afterPropertiesSet() {
        // TODO: 2018/10/10 如果主库为NULL
        if(null == this.masterDataSource){
            throw new IllegalArgumentException("Property 'masterDataSource' is required");
        }
        // TODO: 2018/10/10 设置主库为默认的数据源
        setDefaultTargetDataSource(masterDataSource);

        Map<Object, Object> targetDataSources = new HashMap<>();
        targetDataSources.put(DBType.MASTER.name(), masterDataSource);
        if(null == this.slaveDataSources){
            slaveDataSourceSize = 0;
        }else{
            for(int i=0, len = slaveDataSources.size(); i < len; i++){
                targetDataSources.put(DBType.SLAVE.name() + i, slaveDataSources.get(i));
            }
            slaveDataSourceSize = slaveDataSources.size();
        }

        setTargetDataSources(targetDataSources);
        super.afterPropertiesSet();
    }

    @Nullable
    @Override
    protected Object determineCurrentLookupKey() {
        DBType dbType = DBHolder.getDataSource();
        if(null == dbType || dbType == DBType.MASTER || slaveDataSourceSize <= 0){
            return DBType.MASTER.name();
        }

        int index = 1;

        if(1 == readDataSourcePollPattern){
            //轮询方式
            long currentValue = counter.incrementAndGet();
            if((currentValue + 1) >= MAX_POOL){
                try {
                    lock.lock();
                    if((currentValue + 1) >= MAX_POOL){
                        counter.set(0);
                    }
                } finally {
                    lock.unlock();
                }
            }
            index = (int) (currentValue % slaveDataSourceSize);
        } else {
            //随机方式
            index = ThreadLocalRandom.current().nextInt(0, slaveDataSourceSize);
        }

        return dbType.name() + index;
    }

    public void setMasterDataSource(Object masterDataSource) {
        this.masterDataSource = masterDataSource;
    }

    public void setSlaveDataSources(List<Object> slaveDataSources) {
        this.slaveDataSources = slaveDataSources;
    }

    public void setReadDataSourcePollPattern(int readDataSourcePollPattern) {
        this.readDataSourcePollPattern = readDataSourcePollPattern;
    }
}
